from transformers import AutoTokenizer, AutoModelForCausalLM
from peft import PeftModel
import torch

# 模型路径
model_name_or_path = "./Qwen2.5-7B-Instruct"  # 原始模型路径
lora_model_path = "./qwen-liucixin-lora"  # LoRA 模型路径

# 加载分词器
tokenizer = AutoTokenizer.from_pretrained(model_name_or_path, trust_remote_code=True)
if tokenizer.pad_token is None:
    tokenizer.pad_token = tokenizer.eos_token

# 加载模型
model = AutoModelForCausalLM.from_pretrained(
    model_name_or_path,
    device_map="auto",
    torch_dtype=torch.float16,  # FP16 加速
    trust_remote_code=True
)

# 加载 LoRA 适配器
model = PeftModel.from_pretrained(model, lora_model_path)
model.eval()  # 设置为评估模式
model.to("cuda")  # 将整个模型移动到 GPU，确保与输入张量在同一设备上

# 循环接受用户输入
while True:
    prompt = input("请输入你的提示 (输入 'quit' 退出): ")
    if prompt.lower() == 'quit':
        break

    # 准备输入
    inputs = tokenizer(prompt, return_tensors="pt", padding=True).to("cuda")  # 将输入转换为 PyTorch 张量并移到 GPU

    # 生成文本
    with torch.no_grad():  # 禁用梯度计算，节省显存
        outputs = model.generate(
            **inputs,
            max_length=200,  # 生成文本的最大长度
            do_sample=True,  # 是否使用采样。 如果为 False, 则使用 greedy decoding
            temperature=0.8,  # 控制随机性的参数。 值越高，结果越随机。
            top_p=0.95,  # 与 temperature 结合使用，用于限制采样范围
            top_k=40,  # 与 temperature 结合使用，用于限制采样范围
            repetition_penalty=1.1  # 防止模型重复生成相同的内容。根据情况调整。
        )

    # 解码输出
    generated_text = tokenizer.decode(outputs[0], skip_special_tokens=True)  # 将模型输出的 token IDs 转换为文本
    print(f"生成的文本: {generated_text}")  # 打印生成的文本

print("程序已退出。")
